from PIL import Image, ImageEnhance, ImageFilter, ImageOps
import pytesseract

# 图像增强函数
def enhance_image(image):
    # 转换为灰度图像
    image = image.convert('L')
    # 自适应二值化（局部阈值处理）
    image = ImageOps.autocontrast(image, cutoff=5)
    # 增强对比度
    enhancer = ImageEnhance.Contrast(image)
    image = enhancer.enhance(2.0)  # 对比度增强倍数，可根据需要调整
    # 去噪
    image = image.filter(ImageFilter.MedianFilter(size=3))  # 中值滤波去噪
    return image

# 加载图像
image_path = 'output.png'  # 替换为你的图像路径
img = Image.open(image_path)

# 图像增强
enhanced_image = enhance_image(img)
enhanced_image.show()  # 显示增强后的图像，检查预处理效果

# Tesseract 配置
custom_config = r'--oem 3 --psm 6 -c preserve_interword_spaces=1'  # OCR 引擎模式为 3（LSTM），页面分割模式为 6（假设为一块统一的文本块）

# 使用 image_to_data 获取详细数据
data = pytesseract.image_to_data(enhanced_image, lang='chi_sim', output_type=pytesseract.Output.DICT, config=custom_config)

# 使用 image_to_string 获取完整文本
text = pytesseract.image_to_string(enhanced_image, lang='chi_sim', config=custom_config)

# 检查 data 中的文本是否完整
if not any(data['text']):  # 如果 data 中的文本为空
    print("image_to_data 未识别到文本，使用 image_to_string 的结果")
    data['text'] = text.split('\n')  # 将 text 按行分割，模拟 data 的结构
else:
    print("image_to_data 识别到文本")

# 输出结果
print("image_to_data 结果：")
for i in range(len(data['text'])):
    if data['text'][i].strip():  # 只输出非空文本
        print(f"文本: {data['text'][i]}, 置信度: {data['conf'][i]}, 位置: ({data['left'][i]}, {data['top'][i]})")

print("\nimage_to_string 结果：")
print(text)